import os
import glob
import numpy as np
import trimesh
import trimesh.creation
from environment.unit.Part.Base import Base
from environment.unit.base.Shape import Box
from trimesh.collision import CollisionManager
from trimesh.transformations import transform_points, rotation_matrix, translation_matrix


class Artifacts(object):
    def __init__(self, path='', base_mode=1):
        self.path = path
        self.locating_point_list = []
        self.artifact_base = Base_Env(base_mode, [0, 1200, 0])
        self.transform = self.artifact_base.matrix
        self.transform_inv = self.artifact_base.matrix_inv
        self.env = Env(path)


class Env(object):
    def __init__(self, env_mesh):
        if isinstance(env_mesh, Artifacts):
            env_mesh = trimesh.load(env_mesh.path)
        if isinstance(env_mesh, str):
            env_mesh = trimesh.load(env_mesh)
        self.env_mesh = env_mesh
        self.env_mesh.visual.face_colors = [200, 200, 100, 150]
        self.welds = []

    def add_transform(self, mat):
        self.env_mesh.apply_transform(mat)

    def add_weld(self, origin, direction):
        origin = np.asarray(origin, float)
        direction = np.asarray(direction, float)
        self.welds.append((origin, direction))

    def get_meshes(self):
        meshes = list()
        meshes.append(("env", self.env_mesh))
        for idx, weld in enumerate(self.welds):
            cylinder = trimesh.creation.cylinder(radius=10, height=600)
            cylinder.visual.face_colors = [200, 0, 0, 150]
            cylinder.apply_transform(
                translation_matrix(weld[0])
            )
            cylinder.apply_transform(
                translation_matrix([0, 0, 200])
            )
            meshes.append(("weld" + str(idx), cylinder))
        return meshes

    def get_cm(self):
        env_cm = CollisionManager()
        meshes = self.get_meshes()
        for name, mesh in meshes:
            env_cm.add_object(name, mesh)
        return env_cm


class Base_Env(object):
    def __init__(self, xyz_mode, base_origin=None, size=None):
        if size is None:
            size = [2000, 2000]
        if base_origin is None:
            base_origin = [0, 0, 0]
        self.size = size
        self.xyz_mode = xyz_mode
        self.xyz_normal = 2
        self.dis = 0
        self.base_origin = base_origin
        self.transform = [0, 0, 0]
        self.matrix = None
        self.matrix_inv = None
        self.base = None
        self.base_mesh = None
        self.normal = []
        self.construct()

    def construct(self):
        defaults = {
            'length': self.size[0],
            'width': self.size[1],
            'xyz_mode': self.xyz_mode
        }
        self.base = Base(defaults)

        # self.transform[self.xyz_mode] = 1
        if self.xyz_mode == 0:  # XY
            self.xyz_normal = 2
            self.transform = [1, 1, 1]
            # self.base_origin = [0, 0, self.dis]
            self.dis = self.base_origin[2 - self.xyz_mode]
            self.normal = [0, 0, 1]

            position_point = self.base.get_point('box.-z')
        elif self.xyz_mode == 1:  # XZ
            self.xyz_normal = 1
            self.transform = [1, 0, 0]
            # self.base_origin = [0, self.dis, 0]
            self.dis = self.base_origin[2 - self.xyz_mode]

            position_point = self.base.get_point('box.-y')
            self.normal = [0, 1, 0]

        else:  # YZ
            self.xyz_normal = 0
            self.transform = [0, 1, 0]
            # self.base_origin = [self.dis, 0, 0]
            self.dis = self.base_origin[2 - self.xyz_mode]

            position_point = self.base.get_point('box.-x')
            self.normal = [1, 0, 0]

        self.base.move_point2point(position_point, self.base_origin)
        self.base_mesh = self.base.get_mesh()
        if self.base_mesh.vertices[0][self.xyz_normal] <= self.dis:
            self.matrix = rotation_matrix(
                np.deg2rad(-90),
                self.transform,
                [0, 0, 0]
            )
        else:
            self.matrix = rotation_matrix(
                np.deg2rad(90),
                self.transform,
                [0, 0, 0]
            )
        matrix_mat = np.mat(self.matrix)
        matrix_mat_a = matrix_mat.I
        self.matrix_inv = matrix_mat_a.A

    def add_transform(self, mat):
        self.base.add_transform(mat)
        self.base_mesh.apply_transform(mat)
